#!/usr/bin/env python3
"""
离线日志回放并导出白/黑无人艇轨迹图

用法示例：
  python sh15/tools/log_replay_plot.py --log sh15/logs/方案3_log.json --out sh15/logs/方案3_traj.png
  python sh15/tools/log_replay_plot.py --log sh15/logs/方案10_log.jsonl --use-jsonl --out sh15/logs/方案10_traj.png
"""
import os
import re
import json
import math
import argparse
from typing import Dict, List, Tuple, Any

# 兼容从 tools/ 下直接运行，补充父目录到 sys.path，确保能导入 core.utils
import sys
CURRENT_DIR = os.path.dirname(__file__)
PARENT_DIR = os.path.abspath(os.path.join(CURRENT_DIR, '..'))
if PARENT_DIR not in sys.path:
    sys.path.append(PARENT_DIR)

# 使用无界面后端，适合服务器/命令行环境
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.lines import Line2D
from matplotlib import font_manager as fm

from core.utils import (
    POLYGON_POINTS,
    BREAKTHROUGH_LINE_START,
    BREAKTHROUGH_LINE_END,
    DEPLOYMENT_LINE_START,
    DEPLOYMENT_LINE_END,
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description='离线回放日志并导出白/黑无人艇轨迹图')
    parser.add_argument('--log', required=True, help='日志文件路径（聚合 JSON 或 JSONL）')
    parser.add_argument('--use-jsonl', action='store_true', help='将输入当作 JSONL 流解析')
    parser.add_argument('--out', default=None, help='输出图片路径，默认与日志同名同目录、扩展名为 .png')
    parser.add_argument('--dpi', type=int, default=150, help='输出图片 DPI，默认 150')
    parser.add_argument('--size', type=str, default='1920x1080', help='画布像素尺寸，如 1920x1080')
    parser.add_argument('--with-labels', action='store_true', default=True, help='是否标注终点的 unit_id，默认开启')
    parser.add_argument('--no-labels', dest='with_labels', action='store_false', help='关闭标注 unit_id')
    parser.add_argument('--only-usv', action='store_true', default=True, help='仅绘制无人艇（忽略无人机），默认开启')
    parser.add_argument('--include-uav', dest='only_usv', action='store_false', help='同时绘制无人机（不推荐）')
    return parser.parse_args()


def ensure_chinese_font():
    candidates = [
        # 仓库根目录
        os.path.abspath(os.path.join(CURRENT_DIR, '..', '..', 'NotoSansSC-Regular.ttf')),
        # 工作目录
        os.path.abspath('NotoSansSC-Regular.ttf'),
        # tools 同级（不太可能）
        os.path.abspath(os.path.join(CURRENT_DIR, 'NotoSansSC-Regular.ttf')),
    ]
    for path in candidates:
        if os.path.exists(path):
            try:
                fm.fontManager.addfont(path)
                plt.rcParams['font.sans-serif'] = ['Noto Sans SC']
                plt.rcParams['axes.unicode_minus'] = False
                return
            except Exception:
                pass
    # 回退默认字体
    plt.rcParams['axes.unicode_minus'] = False


def parse_size(size_str: str, dpi: int) -> Tuple[float, float]:
    try:
        w_str, h_str = re.split(r'[xX]', size_str)
        width_px = int(w_str)
        height_px = int(h_str)
        return width_px / dpi, height_px / dpi
    except Exception:
        # 回退 1920x1080
        return 1920 / dpi, 1080 / dpi


def load_aggregated_json(path: str) -> Dict[str, Dict[str, Any]]:
    with open(path, 'r', encoding='utf-8') as f:
        data = json.load(f)
    # 确保键为字符串
    return {str(k): v for k, v in data.items()}


def load_jsonl_as_aggregated(path: str) -> Dict[str, Dict[str, Any]]:
    aggregated: Dict[str, Dict[str, Any]] = {}
    with open(path, 'r', encoding='utf-8') as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                obj = json.loads(line)
            except Exception:
                continue
            t = str(obj.get('time', ''))
            if not t:
                continue
            aggregated[t] = {
                'white': obj.get('white', {}),
                'black': obj.get('black', {}),
            }
    return aggregated


def is_white_usv(name: str) -> bool:
    return bool(re.fullmatch(r'w_usv\d+', name))


def is_black_usv(name: str) -> bool:
    return bool(re.fullmatch(r'b_usv\d+', name))


def is_white_uav(name: str) -> bool:
    return bool(re.fullmatch(r'w_uav\d+', name))


def build_trajectories(agg: Dict[str, Dict[str, Any]], only_usv: bool = True) -> Tuple[Dict[str, List[Tuple[float, float]]], Dict[str, List[Tuple[float, float]]], Tuple[float, float]]:
    # 返回：白方轨迹、黑方轨迹、(t_min, t_max)
    times: List[float] = []
    trajectories_white: Dict[str, List[Tuple[float, float]]] = {}
    trajectories_black: Dict[str, List[Tuple[float, float]]] = {}
    inactive: Dict[str, bool] = {}

    # 按时间排序
    def _to_time_key(s: str) -> float:
        try:
            return float(s)
        except Exception:
            return math.inf

    sorted_times = sorted(agg.keys(), key=_to_time_key)
    if sorted_times:
        times = [_to_time_key(t) for t in sorted_times]
    t_min = times[0] if times else 0.0
    t_max = times[-1] if times else 0.0

    for t in sorted_times:
        frame = agg.get(t, {})
        for side_key, is_white in (('white', True), ('black', False)):
            section = frame.get(side_key, {}) or {}
            for unit_name, unit in section.items():
                # 过滤类型
                if only_usv:
                    if is_white:
                        if not is_white_usv(unit_name):
                            continue
                    else:
                        if not is_black_usv(unit_name):
                            continue
                else:
                    # 若包含 UAV，也允许 w_uav*，但本任务默认 only_usv=True
                    pass

                # 退出后不再追加
                key = f"{side_key}:{unit_name}"
                if inactive.get(key, False):
                    continue

                # 如果当前状态为退出，则本帧不加入点，并标记为 inactive
                state = str(unit.get('state', '')).strip()
                if state == '退出':
                    inactive[key] = True
                    continue

                pos = unit.get('pos', None)
                if not (isinstance(pos, (list, tuple)) and len(pos) >= 2):
                    continue
                try:
                    x = float(pos[0])
                    y = float(pos[1])
                except Exception:
                    continue

                # 写入轨迹
                if is_white:
                    traj = trajectories_white.setdefault(unit_name, [])
                    traj.append((x, y))
                else:
                    traj = trajectories_black.setdefault(unit_name, [])
                    traj.append((x, y))

    return trajectories_white, trajectories_black, (t_min, t_max)


def compute_world_bounds(margin: float = 10000.0) -> Tuple[float, float, float, float]:
    xs = [p[0] for p in POLYGON_POINTS]
    ys = [p[1] for p in POLYGON_POINTS]
    min_x, max_x = min(xs), max(xs)
    min_y, max_y = min(ys), max(ys)
    return (min_x - margin, max_x + margin, min_y - margin, max_y + margin)


def draw_background(ax):
    min_x, max_x, min_y, max_y = compute_world_bounds()
    ax.set_xlim(min_x, max_x)
    ax.set_ylim(min_y, max_y)
    ax.set_aspect('equal', adjustable='box')

    # 背景颜色
    ax.set_facecolor((0.95, 0.96, 0.98))

    # 战区边界
    poly_x = [p[0] for p in POLYGON_POINTS] + [POLYGON_POINTS[0][0]]
    poly_y = [p[1] for p in POLYGON_POINTS] + [POLYGON_POINTS[0][1]]
    ax.plot(poly_x, poly_y, color='white', linewidth=2.0, zorder=1)

    # 突破线 A1-A6（红）
    ax.plot([BREAKTHROUGH_LINE_START[0], BREAKTHROUGH_LINE_END[0]],
            [BREAKTHROUGH_LINE_START[1], BREAKTHROUGH_LINE_END[1]],
            color='red', linewidth=1.5, zorder=1)

    # 部署线 A7-A8（绿）
    ax.plot([DEPLOYMENT_LINE_START[0], DEPLOYMENT_LINE_END[0]],
            [DEPLOYMENT_LINE_START[1], DEPLOYMENT_LINE_END[1]],
            color='green', linewidth=1.5, zorder=1)

    # 轴标签
    ax.set_xlabel('X (m)')
    ax.set_ylabel('Y (m)')


def plot_trajectories(ax,
                      white_traj: Dict[str, List[Tuple[float, float]]],
                      black_traj: Dict[str, List[Tuple[float, float]]],
                      with_labels: bool = True):
    # 颜色方案
    white_line_color = '#6495ED'  # cornflower blue
    white_point_color = '#A7C7FF'
    black_line_color = '#FF7F7F'
    black_point_color = '#FFC0C0'

    # 白方
    for unit_name, points in white_traj.items():
        if len(points) < 1:
            continue
        xs = [p[0] for p in points]
        ys = [p[1] for p in points]
        ax.plot(xs, ys, color=white_line_color, linewidth=1.8, alpha=0.9, zorder=2)
        ax.scatter(xs[::max(1, len(xs)//200)], ys[::max(1, len(ys)//200)], s=6, color=white_point_color, alpha=0.6, zorder=3)
        # 起止标记
        start_x, start_y = points[0]
        end_x, end_y = points[-1]
        ax.scatter([start_x], [start_y], marker='^', s=40, color='green', zorder=4)
        ax.scatter([end_x], [end_y], marker='s', s=40, color='orange', zorder=4)
        if with_labels:
            try:
                idx = int(re.findall(r'\d+', unit_name)[0])
                label = f"W{idx}"
            except Exception:
                label = f"W:{unit_name}"
            ax.text(end_x + 500, end_y + 500, label, fontsize=8, color=white_line_color)

    # 黑方
    for unit_name, points in black_traj.items():
        if len(points) < 1:
            continue
        xs = [p[0] for p in points]
        ys = [p[1] for p in points]
        ax.plot(xs, ys, color=black_line_color, linewidth=1.8, alpha=0.9, zorder=2)
        ax.scatter(xs[::max(1, len(xs)//200)], ys[::max(1, len(ys)//200)], s=6, color=black_point_color, alpha=0.6, zorder=3)
        # 起止标记
        start_x, start_y = points[0]
        end_x, end_y = points[-1]
        ax.scatter([start_x], [start_y], marker='^', s=40, color='green', zorder=4)
        ax.scatter([end_x], [end_y], marker='s', s=40, color='orange', zorder=4)
        if with_labels:
            try:
                idx = int(re.findall(r'\d+', unit_name)[0])
                label = f"B{idx}"
            except Exception:
                label = f"B:{unit_name}"
            ax.text(end_x + 500, end_y + 500, label, fontsize=8, color=black_line_color)

    # 简洁图例（不对每条轨迹逐一列出）
    legend_elements = [
        Line2D([0], [0], color='#6495ED', lw=2, label='白方无人艇轨迹'),
        Line2D([0], [0], color='#FF7F7F', lw=2, label='黑方无人艇轨迹'),
        Line2D([0], [0], marker='^', color='w', markerfacecolor='green', markersize=8, label='起点'),
        Line2D([0], [0], marker='s', color='w', markerfacecolor='orange', markersize=8, label='终点'),
    ]
    ax.legend(handles=legend_elements, loc='upper right')


def main():
    args = parse_args()
    ensure_chinese_font()

    log_path = args.log
    if not os.path.exists(log_path):
        raise FileNotFoundError(f'日志文件不存在: {log_path}')

    if args.out:
        out_path = args.out
    else:
        base = os.path.splitext(log_path)[0]
        out_path = base + '_traj.png'
    os.makedirs(os.path.dirname(out_path) or '.', exist_ok=True)

    # 读取日志
    aggregated = load_jsonl_as_aggregated(log_path) if args.use_jsonl else load_aggregated_json(log_path)

    # 构建轨迹
    white_traj, black_traj, (t_min, t_max) = build_trajectories(aggregated, only_usv=args.only_usv)

    # 画布
    figsize = parse_size(args.size, args.dpi)
    fig, ax = plt.subplots(figsize=figsize, dpi=args.dpi)
    draw_background(ax)
    plot_trajectories(ax, white_traj, black_traj, with_labels=args.with_labels)

    # 标题
    file_name = os.path.basename(log_path)
    duration = max(0, int(round(t_max - t_min)))
    title = f"轨迹回放: {file_name} | 时长: {duration}秒"
    ax.set_title(title)

    plt.tight_layout()
    fig.savefig(out_path)
    plt.close(fig)
    print(f"已导出: {out_path}")


if __name__ == '__main__':
    main()


